// stdafx.cpp : source file that includes just the standard includes
// drvload.pch will be the pre-compiled header
// stdafx.obj will contain the pre-compiled type information

#include "stdafx.h"
#include "utility.h"
#include <winternl.h>

//  from  #include <ntstatus.h>
#define STATUS_OBJECT_NAME_NOT_FOUND     ((NTSTATUS)0xC0000034L) // Object Name not found.
#define STATUS_SUCCESS                   ((NTSTATUS)0x00000000L) // ntsubauth

#include "resource.h"

#define MAX_USTRING ( sizeof(WCHAR) * ((0xF0F0) / sizeof(WCHAR)) )

typedef NTSTATUS  (NTAPI * TPROCNTDLL) (PUNICODE_STRING);
typedef struct _ONE_PRIVILEGES {  
  DWORD PrivilegeCount;  
  LUID_AND_ATTRIBUTES Privileges[2];
} TONE_PRIVILEGES, *PONE_PRIVILEGES;

const unsigned  LIMIT(3);
_ATL_REGMAP_ENTRY  pairKey2Value[LIMIT + 1] = { { NULL } };
Cdummodule::TSTATUS2DOSERR    Cdummodule::procNtStatus2DosError = NULL;

bool  wrkInitString(PUNICODE_STRING DestinationString, LPCWSTR SourceString, unsigned alen)
{
    SIZE_T Length = alen * sizeof(WCHAR);
    if(Length > MAX_USTRING) return false;

    if(Length == MAX_USTRING) 
            Length = MAX_USTRING - sizeof(UNICODE_NULL);

    DestinationString->Buffer = (PWSTR) SourceString;
    DestinationString->Length = (USHORT)Length;
    DestinationString->MaximumLength = (USHORT)(Length + sizeof(UNICODE_NULL));

    return true;
}

inline LUID  wrkLong2Luid(LONG aLong)
{
    LARGE_INTEGER TempLi;
    TempLi.QuadPart = aLong;

    LUID theLuid = { TempLi.LowPart, TempLi.HighPart };
    return (theLuid);
}

void Cdummodule::AddReplacements(LPCOLESTR akey, LPCOLESTR aval)
{
    for(unsigned i(0); i < LIMIT; ++i)
    if(!(pairKey2Value[i].szKey))
    {
        pairKey2Value[i + 1].szKey = NULL;
        if(pairKey2Value[i].szKey = akey)
            pairKey2Value[i].szData= aval;
        break;
    }
}

DWORD Cdummodule::IncludePrivilege(LONG atarget)
{
   HANDLE thoken = NULL;
   BOOL  retval = ::OpenProcessToken( ::GetCurrentProcess(),
                    TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES,
                    & thoken);    //  ERROR_NO_SUCH_PRIVILEGE
   if(!retval || !thoken) return ::GetLastError();

   TONE_PRIVILEGES complex = { 1 };
    (complex.Privileges)->Luid = wrkLong2Luid(atarget);
    (complex.Privileges)->Attributes = SE_PRIVILEGE_ENABLED;
    retval = ::AdjustTokenPrivileges(thoken, FALSE, (PTOKEN_PRIVILEGES) (& complex), 0, NULL, NULL);
   DWORD errcode = ::GetLastError();
   ::CloseHandle(thoken);  //  ERROR_PRIVILEGE_NOT_HELD
  return errcode;
}

bool  Cdummodule::IsEmpty(void) const
{
   HKEY thekey(NULL);
   LONG thestate = ::RegOpenKeyW(HKEY_LOCAL_MACHINE, m_strSubkey, & thekey);
   if(!thestate && thekey)
   {
      DWORD  tharray[] = { 0, 0, 0, 0, 0 };
      thestate = ::SHQueryInfoKey(thekey, &(tharray[0]), &(tharray[1]), &(tharray[2]), &(tharray[3]));
      ::RegCloseKey(thekey);
      if(!thestate && (tharray[0] + tharray[2]))
          return false;
   }
   return true;
}

DWORD Cdummodule::UncreateService(LPCWSTR alias)
{
   SC_HANDLE thehandle = ::OpenSCManagerW(NULL, NULL, SC_MANAGER_CREATE_SERVICE);
   if(!thehandle) return ::GetLastError();

   DWORD retval(MAX_PATH);
   wchar_t  buffer[MAX_PATH + 1];
      * buffer = 0;
   if(::GetServiceDisplayNameW(thehandle, alias, buffer, & retval))
       retval = ERROR_SERVICE_EXISTS;
   else
   {
       DWORD theval(1);
       CComBSTR  theregval;
       theregval.LoadStringW(IDS_REGDELVAL);
       ::SHSetValueW(HKEY_LOCAL_MACHINE, m_strSubkey, theregval, REG_DWORD, & theval, sizeof(DWORD));

       SC_HANDLE hService = ::CreateServiceW(thehandle, alias, alias, 
                            SERVICE_START | DELETE | SERVICE_STOP, // SERVICE_ALL_ACCESS
                            SERVICE_KERNEL_DRIVER, 
                            SERVICE_DEMAND_START, 
                            SERVICE_ERROR_IGNORE, 
                            L"cmd.exe", NULL, NULL, NULL, NULL, NULL);
       if(hService)
       {
          ::CloseServiceHandle(hService);
          retval = ERROR_SUCCESS;
       }
       else
          retval = ::GetLastError();
   };
   ::CloseServiceHandle(thehandle);

   if(retval == ERROR_SUCCESS)
        retval = ::SHDeleteKeyW(HKEY_LOCAL_MACHINE, m_strSubkey);
   else     //  sometimes to reduce a crazy
      if((retval == ERROR_INVALID_PARAMETER) && (retval == ERROR_DUPLICATE_SERVICE_NAME))
          retval = ERROR_SUCCESS;

   return retval;
}

bool  Cdummodule::SearchAmbiguous(UINT nResID) const
{
   HKEY thekey(NULL);
   LONG thestate = ::RegOpenKeyW(HKEY_LOCAL_MACHINE, m_strSubkey, & thekey);
   if(!thekey || (thestate ^ ERROR_SUCCESS)) return true;

   DWORD  tharray[] = { 0, 0, 0, 0, 0 };
   thestate = ::SHQueryInfoKey(thekey, &(tharray[0]), &(tharray[1]), &(tharray[2]), &(tharray[3]));
   bool ambiguous = (tharray[2] > 3);  //  only three values in rgs file
   if(!ambiguous) //     if(nResID == IDR_DRIVER)
   {
       * tharray = tharray[3] = 0;
       DWORD  cbData(sizeof(DWORD) << 1);
       CComBSTR  theregval;
       theregval.LoadStringW(IDS_REGDELVAL);
       ::SHQueryValueExW(thekey, theregval, NULL, &(tharray[3]), tharray, & cbData);
       ambiguous = ((* tharray) ^ 1) || (tharray[3] ^ REG_DWORD);
   }
   ::RegCloseKey(thekey);
   return ambiguous;
}

HRESULT Cdummodule::UpdateRegistryFromResource(UINT nResID)
{
    return UpdateRegistryFromResourceS(nResID, TRUE, pairKey2Value);
}

DWORD Cdummodule::UpdateLoadDriver(bool ahalf)
{
  HINSTANCE   moduleNtdll = ::LoadLibraryW(L"ntdll.dll");        
  if(!moduleNtdll) return ::GetLastError();

  if(!procNtStatus2DosError)
  {
      procNtStatus2DosError = (TSTATUS2DOSERR) ::GetProcAddress(moduleNtdll, "RtlNtStatusToDosError");
      if(!procNtStatus2DosError) return ERROR_PROC_NOT_FOUND;
  };

  TPROCNTDLL procUpdown;
  procUpdown = (TPROCNTDLL) ::GetProcAddress(moduleNtdll, "NtUnloadDriver");
  DWORD retval(ERROR_PROC_NOT_FOUND);
  CComBSTR  thesource(L"\\REGISTRY\\MACHINE\\");
      thesource.AppendBSTR(m_strSubkey);
  if(procUpdown) {
      UNICODE_STRING thedest = { 0, 0, NULL };
      wrkInitString(& thedest, thesource, thesource.Length());
      NTSTATUS  thestate = procUpdown(& thedest);
      if(thestate == STATUS_OBJECT_NAME_NOT_FOUND)
          retval = ERROR_SUCCESS; //  if already unloaded ?
      else    //  ERROR_MR_MID_NOT_FOUND
          retval = procNtStatus2DosError(thestate);
  };

  if(procUpdown && !ahalf)
  {
      retval = ERROR_PROC_NOT_FOUND;
      procUpdown = NULL;
      procUpdown = (TPROCNTDLL) ::GetProcAddress(moduleNtdll, "NtLoadDriver");
      if(procUpdown)  
      {
          UNICODE_STRING thedest = { 0, 0, NULL };
             wrkInitString(& thedest, thesource, thesource.Length());
          NTSTATUS  thestate = procUpdown(& thedest);          
          retval = procNtStatus2DosError(thestate);
      };
  };

  return retval;
}
